import numpy as np
import cv2
import random

def generate_random_nv12(width=100, height=100):
    """
    生成随机NV12格式的图片数据
    """
    y_component = np.random.randint(0, 256, (height, width), dtype=np.uint8)
    uv_height = height // 2
    uv_component = np.random.randint(0, 256, (uv_height, width), dtype=np.uint8)
    return np.vstack((y_component, uv_component))

def read_nv12_from_file(file_path, width=640, height=480):
    """
    从文件读取NV12格式的图片数据
    """
    expected_size = int(height * width * 3 / 2)
    with open(file_path, 'rb') as f:
        raw_data = f.read()
    if len(raw_data) != expected_size:
        raise ValueError(f"文件尺寸不匹配，预期 {expected_size} 字节，实际 {len(raw_data)} 字节")
    uv_height = height // 2
    return np.frombuffer(raw_data, dtype=np.uint8).reshape((height + uv_height, width))

def nv12_to_rgb(nv12_data, width=None, height=None):
    """
    将NV12格式的数据转换为RGB格式（不依赖OpenCV）
    """
    if width is None or height is None:
        total_height = nv12_data.shape[0]
        height = int(total_height * 2 / 3)
        width = nv12_data.shape[1]
    
    y = nv12_data[:height, :]
    uv = nv12_data[height:, :].reshape((-1, width))
    
    u = uv[:, 0::2].astype(np.float32) - 128
    v = uv[:, 1::2].astype(np.float32) - 128
    
    u = np.repeat(np.repeat(u, 2, axis=0), 2, axis=1)[:height, :width]
    v = np.repeat(np.repeat(v, 2, axis=0), 2, axis=1)[:height, :width]
    
    y = y.astype(np.float32)
    r = np.clip(y + 1.402 * v, 0, 255)
    g = np.clip(y - 0.344136 * u - 0.714136 * v, 0, 255)
    b = np.clip(y + 1.772 * u, 0, 255)
    
    return np.stack([r, g, b], axis=-1).astype(np.uint8)

def generate_random_crop_region(img_width, img_height, crop_size=(100, 100)):
    """
    生成随机裁剪区域
    
    :param img_width: 原图宽度
    :param img_height: 原图高度
    :param crop_size: 裁剪区域大小 (width, height)
    :return: (x, y, w, h) 裁剪区域坐标和大小
    """
    crop_w, crop_h = crop_size
    max_x = img_width - crop_w
    max_y = img_height - crop_h
    
    if max_x < 0 or max_y < 0:
        raise ValueError("裁剪尺寸大于原图尺寸")
    
    x = random.randint(0, max_x)
    y = random.randint(0, max_y)
    
    return (x, y, crop_w, crop_h)

def crop_nv12(nv12_data, crop_region):
    """
    裁剪NV12图像（支持任意尺寸），自动调整x,y为偶数以保证UV分量对齐
    
    :param nv12_data: NV12格式图像数据
    :param crop_region: (x, y, w, h) 裁剪区域
    :return: 裁剪后的NV12图像
    """
    x, y, w, h = crop_region
    height_total = nv12_data.shape[0]
    
    # 计算原始Y分量的高度（占图像总高度的2/3）
    y_height = height_total * 2 // 3
    
    # 确保裁剪区域在有效范围内
    orig_width = nv12_data.shape[1]
    orig_height_y = y_height
    
    # 调整x,y为偶数
    x = max(0, min(x & ~1, orig_width - 1))  # 向下取最近的偶数
    y = max(0, min(y & ~1, orig_height_y - 1))  # 向下取最近的偶数
    
    # 调整宽度和高度，确保不超过边界
    w = max(2, min(w, orig_width - x))  # 最小宽度为2，保证UV分量完整
    h = max(2, min(h, orig_height_y - y))  # 最小高度为2，保证UV分量完整
    
    # 确保宽度和高度也是偶数
    w = w & ~1
    h = h & ~1
    
    # 裁剪Y分量
    cropped_y = nv12_data[y:y+h, x:x+w]
    
    # 计算UV分量的起始位置和高度（NV12格式中UV高度是Y高度的一半）
    uv_start = y_height
    uv_height_total = height_total - uv_start
    uv_component = nv12_data[uv_start:]
    
    # UV垂直方向的偏移（y/2）
    uv_y = y // 2
    
    # UV垂直高度（h/2）
    uv_h = h // 2
    
    # 确保UV位置不越界
    uv_y = max(0, min(uv_y, uv_height_total - 1))
    uv_h = max(1, min(uv_h, uv_height_total - uv_y))
    
    # 直接使用相同的x和w进行裁剪（x已经是偶数，w也是偶数）
    cropped_uv = uv_component[uv_y:uv_y+uv_h, x:x+w]
    
    # 合并裁剪后的Y和UV
    return np.vstack((cropped_y, cropped_uv))


def visualize_nv12(nv12_data, crop_region=None, save_path=None):
    """
    可视化NV12图像（转换为RGB显示），并可选择保存图像
    如果提供了裁剪区域，会在图像上绘制裁剪框
    
    参数:
        nv12_data: NV12格式的图像数据
        crop_region: (可选)裁剪区域 (x, y, w, h)
        save_path: (可选)如果提供，将保存图像到此路径
    """
    rgb_image = nv12_to_rgb(nv12_data)
    bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
    
    # 如果提供了裁剪区域，绘制矩形框
    if crop_region is not None:
        x, y, w, h = crop_region
        # 绘制红色矩形框，线宽2
        cv2.rectangle(bgr_image, (x, y), (x+w, y+h), (0, 0, 255), 2)
        # 创建可调整大小的窗口
    cv2.namedWindow('NV12 Image', cv2.WINDOW_NORMAL)
    # 设置窗口大小
    cv2.resizeWindow('NV12 Image', 720, 640)
    # 显示图像
    cv2.imshow('NV12 Image', bgr_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    
    # 如果需要保存图像
    if save_path is not None:
        cv2.imwrite(save_path, bgr_image)
        print(f"图像已保存到: {save_path}")



if __name__ == "__main__":
    # 生成随机NV12图像
    width, height = 1920, 1080
    # nv12_data = generate_random_nv12(width, height)
    nv12_data = read_nv12_from_file("test_w1920_h1080.nv12", width,height )
    # # 生成随机裁剪区域
    min_crop_size = 100
    max_crop_size = min(width, height) // 2
    crop_w = random.randint(min_crop_size, max_crop_size)
    crop_h = random.randint(min_crop_size, max_crop_size)
    crop_size = (crop_w, crop_h)
    
    crop_region = generate_random_crop_region(width, height, crop_size)
    print(f"裁剪区域: {crop_region[0]}, {crop_region[1]}, {crop_region[2]}, {crop_region[3]}")

    # 裁剪图像
    # crop_region = (280, 330, 501, 501)
    cropped_nv12 = crop_nv12(nv12_data, crop_region)
    
    # # 显示原始图像和裁剪后的图像
    # print("原始图像:")
    visualize_nv12(nv12_data=nv12_data, crop_region=crop_region, save_path="test/all.jpg")
    
    # print("裁剪后的图像:")
    # visualize_nv12(nv12_data=cropped_nv12, crop_region=None, save_path="crop.jpg")